import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules
import matplotlib.pyplot as plt
import networkx as nx

# 1. 加载数据
df = pd.read_excel('C:\\Users\\Administrator\\Desktop\\关联规则\\餐厅数据.xlsx') 

# 2. 转换菜品格式：将逗号分隔的字符串拆分为列表
transactions = df['菜品'].str.split(',').to_list()

# 3. 数据标准化：转换为布尔矩阵（One-Hot编码）
te = TransactionEncoder()
te_ary = te.fit(transactions).transform(transactions)
df_encoded = pd.DataFrame(te_ary, columns=te.columns_)

# 4. 使用Apriori算法找出频繁项集
frequent_itemsets = apriori(df_encoded, min_support=0.1, use_colnames=True)
frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)

# 5. 打印二项集
print(frequent_itemsets[frequent_itemsets.itemsets.apply(lambda x: len(x) == 2)])

# 6. 生成关联规则
rules = association_rules(frequent_itemsets, metric="lift", min_threshold=1)  # 新增此行

# 7. 筛选有效规则（提升度>1且置信度>0.1）
effective = rules[
    (rules['lift'] > 1) & (rules['confidence'] > 0.1)
].sort_values(by=['lift', 'confidence'], ascending=False)

# 8. 打印筛选后的规则
print(effective[["antecedents", "consequents", "support", "confidence", "lift"]])

# ...前面的代码保持不变...

# 可视化部分修正
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

G = nx.DiGraph()
for _, row in effective.iterrows():
    G.add_edge(
        ', '.join(list(row['antecedents'])),
        ', '.join(list(row['consequents'])),
        weight=row['lift']
    )

plt.figure(figsize=(12, 8))
pos = nx.spring_layout(G, k=0.5)
ax = plt.gca()  # 获取当前Axes

nx.draw(G, pos,
        ax=ax,
        with_labels=True,
        edge_color=[G[u][v]['weight'] for u, v in G.edges()],
        width=2.0,
        node_size=800,
        alpha=0.7,
        font_size=10,
        arrowsize=20)

edge_weights = [G[u][v]['weight'] for u, v in G.edges()]
sm = plt.cm.ScalarMappable(cmap=plt.cm.viridis,
                          norm=plt.Normalize(vmin=min(edge_weights),
                                           vmax=max(edge_weights)))
sm.set_array([])  # 必须设置
plt.colorbar(sm, ax=ax, label='关联规则提升度(Lift)')

plt.title("菜品关联规则网络图", fontsize=15)
plt.tight_layout()
plt.show()